import paddle
import paddle.nn as nn
from paddle.nn import Conv2D, Linear, Embedding
import paddle.nn.functional as F
import numpy as np
import os


class TextClassifier:
    def __init__(self, vocab_path, model_path, seq_len=150):
        """
        初始化分类器，加载模型和词典。
        :param vocab_path: 字典文件路径
        :param model_path: 模型参数文件路径
        :param seq_len: 模型的最大序列长度
        """
        # 加载词典
        self.vocab = self.load_vocab(vocab_path)
        self.seq_len = seq_len

        # 初始化模型
        self.model = Conv_FFD(dict_dim=self.vocab["<pad>"])
        self.model.eval()

        # 加载模型参数
        model_state_dict = paddle.load(model_path)
        self.model.set_state_dict(model_state_dict)

        # 分类标签映射
        self.label_map = {0: "谣言", 1: "不是谣言", 2: "尚无定论"}

    def load_vocab(self, file_path):
        """
        加载字典文件。
        :param file_path: 字典文件路径
        :return: 字典对象
        """
        with open(file_path, 'r', encoding="utf8") as fr:
            vocab = eval(fr.read())
        return vocab

    def preprocess_input(self, text):
        """
        对输入文本进行编码和填充。
        :param text: 输入的单个文本字符串
        :return: Paddle 张量
        """
        # 将文本编码为索引
        encoded = [self.vocab.get(char, self.vocab.get("<unk>")) for char in text]
        if len(encoded) < self.seq_len:
            encoded += [self.vocab.get("<pad>")] * (self.seq_len - len(encoded))  # 填充
        else:
            encoded = encoded[:self.seq_len]  # 截断

        # 转换为张量，增加 batch 维度
        input_tensor = paddle.to_tensor([encoded], dtype='int64')  # Shape: [1, seq_len]
        return input_tensor

    def predict(self, input_text):
        """
        对输入文本进行预测。
        :param input_text: 输入文本字符串
        :return: 预测结果字符串和置信度
        """
        # 预处理输入文本
        input_tensor = self.preprocess_input(input_text)

        # 模型预测
        results = self.model(input_tensor)
        results = F.softmax(results, axis=-1).numpy()  # 转化为概率分布

        # 获取预测类别及置信度
        idx = np.argmax(results[0])  # 最大概率的索引
        confidence = results[0][idx]  # 对应置信度
        predicted_label = self.label_map[idx]

        # 返回结果
        return predicted_label, confidence


# 定义卷积模型
class Conv_FFD(nn.Layer):
    def __init__(self, dict_dim):
        super(Conv_FFD, self).__init__()
        self.dict_dim = dict_dim
        self.emb_dim = 128
        self.hid_dim = 128
        self.fc_hid_dim = 96
        self.class_dim = 3
        self.channels = 1
        self.win_size = [3, self.hid_dim]
        self.seq_len = 150
        self.embedding = Embedding(self.dict_dim + 1, self.emb_dim, sparse=False)
        self.hidden1 = Conv2D(
            in_channels=1,
            out_channels=self.hid_dim,
            kernel_size=self.win_size,
            padding=[1, 1]
        )
        self.relu1 = nn.ReLU()
        self.hidden3 = paddle.nn.MaxPool2D(kernel_size=2, stride=2)
        self.hidden4 = Linear(128 * 75, 3)

    def forward(self, input):
        batch_size = input.shape[0]
        x = self.embedding(input)
        x = F.dropout(x, p=0.5)
        x = paddle.reshape(x, [batch_size, 1, self.seq_len, self.emb_dim])
        x = self.hidden1(x)
        x = self.relu1(x)
        x = self.hidden3(x)
        x = paddle.reshape(x, shape=[batch_size, -1])
        out = self.hidden4(x)
        return out


# 示例用法
if __name__ == "__main__":
    here = os.path.dirname(os.path.abspath(__file__))
    vocab_path = os.path.join(here,"data/dict.txt")  # 字典文件路径
    model_path = os.path.join(here, "../src/model_final.pdparams")  # 模型参数文件路径

    # 初始化分类器
    classifier = TextClassifier(vocab_path, model_path)

    # 输入文本
    input_text = input("请输入文本进行预测: ")

    # 预测结果
    label, confidence = classifier.predict(input_text)
    print(f"输入: {input_text}\n预测: {label}\n置信度: {confidence:.2f}")
